#include <number_recog.hpp>
#include <camera.hpp>
#include <video.h>

int test(int COLOR, int tag, string &pathOfDescriptionFile) {
    Video video;
    Mat toBeDetect;
    Mat toBeWritten;
    // const int LOW_EXP = 100;
    const int HIGH_EXP = 2000;
    ofstream os;
    Armor_recog ar(COLOR, 0);
    Number_recog nr;
    os.open(pathOfDescriptionFile, std::ofstream::app);
    // video.videoClose();
    video.videoCheck();
    video.videoOpen();

    video.CameraChangeTrig(Video::ETrigType::trigContinous);
    video.setBalanceRatio(1, 1, 1);
    video.loadSetting(0);
    video.SetExposeTime(HIGH_EXP);

    video.videoStart();
    video.startGrabbing();
    char file_cnt[100];
    int cnt = 0;
    namedWindow("toBeWritten", WINDOW_FREERATIO);
    namedWindow("test", WINDOW_FREERATIO);
    while(true) {
        
        sprintf(file_cnt, "etc/numbers/n%d_%d.png", tag, cnt);
        video.getFrame(toBeDetect);
        // video.SetExposeTime(HIGH_EXP);
        video.getFrame(toBeWritten);
        // video.SetExposeTime(LOW_EXP);
        imshow("toBeWritten", toBeWritten);
        ar.detect_armor(toBeDetect);
        cv::waitKey(1);
        if(ar.armor_detected.size() == 0) continue;
        RotatedRect bound = ar.armor_detected[0].bound;
        float width = ar.armor_detected[0].barLen * ar.ARMOR_HEIGHT_BAR_LENGTH_RATIO;
        bound = RotatedRect(bound.center, Size2f(width, width), bound.angle );
        if(nr.prepare_data(toBeWritten, bound, os, tag)) {
            imwrite(file_cnt, toBeWritten);
            cnt++;
            printf("image of %d No.%d\n", tag, cnt);
        }
    }
}

Number_recog::Number_recog(){
    hog = HOGDescriptor(
        Size(15,15), //winSize
        Size(10,10), //blocksize
        Size(5,5), //blockStride,
        Size(5,5), //cellSize,
                 9, //nbins,
                  1, //derivAper,
                 -1, //winSigma,
        HOGDescriptor::L2Hys, //histogramNormType,
                0.2, //L2HysThresh,
                  1,//gammal correction,
                  64,//nlevels=64
                  1
    );//Use signed gradients
};

bool Number_recog::prepare_data(Mat &src, RotatedRect &bound, ofstream &out_data, int tag)
{
    Mat dst;
    dst = pre_processing(src, bound);
    dst = decrease_to_size(dst, 16);
    imshow("test", dst);
    int key = waitKey(0);
    //pressed d to filter out the unwanted frames
    //other keys writes the data to dataset
    if (key != 'd') {
        write(out_data, dst, tag);
        return true;
    }
    return false;
}

bool Number_recog::read_model(const string &filename_model) {
    this->model = StatModel::load<SVM>(filename_model);
    if (model.empty()) {
        return true;
    }else {
        return false;
    }
}

void Number_recog::train_model(const string &filename_dataset, 
                     const string& filename_model_to_save) 
{
    build_svm_classifier(filename_dataset, filename_model_to_save, string());
}

int Number_recog::predict(const Mat &to_predict, RotatedRect bounding_box) {
    if(model.empty()) {
        cerr << "model not trained or load" << endl;
        return -1;
    }
    
    Mat dst = pre_processing(to_predict, bounding_box);
    dst = decrease_to_size(dst, 16);
    Mat test_data(1, 144, CV_32F);
    namedWindow("dst", WINDOW_FREERATIO);
    imshow("dst", dst);
    // waitKey(0);
    float r;
    std::vector<float> descriptor;
    hog.compute(dst, descriptor);
    int i;
    for(i = 0; i < hog.getDescriptorSize(); i++)
    {
        test_data.at<float>(i) = descriptor[i];
    }
    r = model->predict(test_data);
    return r;
}

cv::Mat Number_recog::pre_processing(const Mat &src, RotatedRect bounding_box) {
    Mat dst = src.clone();
    Point2f bound_src[4];
    Point2f bound_dst[4];
    Point2f _bound_src[4];  //bound_src modified
    Mat warpMat;
    int min_point = 0;
    bounding_box.points(bound_src);
    
    cvtColor(dst, dst, COLOR_RGB2GRAY);
    //get the upper right point of the rotated rectangle
    for(int i = 1; i < 4; i++) {
        if((bound_src[min_point].x + bound_src[min_point].y) > (bound_src[i].x + bound_src[i].y)) {
            min_point = i;
        }
    }

    //let the new bound start with the upper right point, clockwise
    for(int i = 0; i < 4; i++) {
        _bound_src[i] = bound_src[(min_point + i) % 4];
    }
    
    //use perspective Transformation to normalize the src image
    bound_dst[0] = cv::Point2f(0, 0);
    bound_dst[1] = cv::Point2f(256, 0);
    bound_dst[2] = cv::Point2f(256, 256);
    bound_dst[3] = cv::Point2f(0, 256);
    warpMat = getPerspectiveTransform(_bound_src, bound_dst);
    cv::warpPerspective(dst, dst, warpMat,dst.size(), INTER_LINEAR);
    
    Rect bound_rec(bound_dst[0], bound_dst[2]);
    cv::rectangle(dst, bound_rec, Scalar(255, 0, 0), 5);
    threshold(dst, dst, 5, 255, THRESH_BINARY);
    return dst(bound_rec);
}



cv::Mat Number_recog::decrease_to_size(const Mat &src, int size) {
    float temp = log2f(min(src.rows, src.cols) / size);
    int temp1 = temp;
    float temp2 = temp1 + 1 - temp;
    int newSize = size * powf(2, temp1);
    //decrease the size of the photo to size x size
    Mat partial = Mat(src, Rect2i(Point2i((src.cols - newSize) / 2,
                                          (src.rows - newSize) / 2),
                                  Point2i((src.cols + newSize) / 2,
                                          (src.rows + newSize) / 2)));

    for (int i = 0; i < temp1; i++)
    {
        pyrDown(partial, partial, Size2i(partial.cols / 2, partial.rows / 2));
    }
    //threshold(partial, partial, 130, 255, THRESH_BINARY_INV);
    return partial;
}

void Number_recog::write(ofstream &outDataBase, const Mat &toWrite, int tag) {
    outDataBase << tag << ' ';
    std::vector<float> descriptor;
    hog.compute(toWrite, descriptor);
    outDataBase << setprecision(2);
    cout << setprecision(2);
    int i;
    for(i = 0; i < hog.getDescriptorSize(); i++)
    {
        outDataBase << descriptor[i] << ' ';
        std::cout << descriptor[i] << ' ';
    }
    outDataBase << std::endl;
    std::cout << endl;
}
